#include <linux/errno.h>
#include <linux/radix-tree.h>
#include <linux/gfp.h>
#include <linux/rcupdate.h>
#include <linux/types.h>
#include <linux/kernel.h>

#include <asm/non-atomic.h>

#define RADIX_TREE_MAP_SHIFT	(6)

#define RADIX_TREE_MAP_SIZE	(1UL << RADIX_TREE_MAP_SHIFT)
#define RADIX_TREE_MAP_MASK	(RADIX_TREE_MAP_SIZE-1)

#define RADIX_TREE_TAG_LONGS	\
	((RADIX_TREE_MAP_SIZE + BITS_PER_LONG - 1) / BITS_PER_LONG)

struct radix_tree_node {
	unsigned int	height;		/* Height from the bottom */
	unsigned int	count;
	// struct rcu_head	rcu_head;
	void		*slots[RADIX_TREE_MAP_SIZE];
	unsigned long	tags[RADIX_TREE_MAX_TAGS][RADIX_TREE_TAG_LONGS];
};

#define RADIX_TREE_INDEX_BITS  (8 /* CHAR_BIT */ * sizeof(unsigned long))
#define RADIX_TREE_MAX_PATH (DIV_ROUND_UP(RADIX_TREE_INDEX_BITS, \
					  RADIX_TREE_MAP_SHIFT))

static unsigned long height_to_maxindex[RADIX_TREE_MAX_PATH + 1];

static inline void tag_set(struct radix_tree_node *node, unsigned int tag,
		int offset)
{
	__set_bit(offset, node->tags[tag]);
}

static inline int root_tag_get(struct radix_tree_root *root, unsigned int tag)
{
	return (unsigned)root->gfp_mask & (1 << (tag + __GFP_BITS_SHIFT));
}

static inline unsigned long radix_tree_maxindex(unsigned int height)
{
#if 0
	return height_to_maxindex[height];
#else
	switch (height) {
	case 0:
		return 0;
		break;
	case 1:
		return 63;
		break;
	case 2:
		return 4095;
		break;
	case 3:
		return 262143;
		break;
	case 4:
		return 16777215;
		break;
	case 5:
		return 1073741823;
		break;
	case 6:
		return -1;
		break;
	default:
		return -1;
		break;
	}
#endif
}

static struct radix_tree_node *radix_tree_node_alloc(struct radix_tree_root *root)
{
	struct radix_tree_node *ret = NULL;
	// gfp_t gfp_mask = root_gfp_mask(root);

	// if (!(gfp_mask & __GFP_WAIT)) {
	// 	struct radix_tree_preload *rtp;

	// 	/*
	// 	 * Provided the caller has preloaded here, we will always
	// 	 * succeed in getting a node here (and never reach
	// 	 * kmem_cache_alloc)
	// 	 */
	// 	rtp = &__get_cpu_var(radix_tree_preloads);
	// 	if (rtp->nr) {
	// 		ret = rtp->nodes[rtp->nr - 1];
	// 		rtp->nodes[rtp->nr - 1] = NULL;
	// 		rtp->nr--;
	// 	}
	// }
	if (ret == NULL)
		// ret = kmem_cache_alloc(radix_tree_node_cachep, gfp_mask);
		ret = kmalloc(sizeof(struct radix_tree_node));
		memset(ret, 0, sizeof(struct radix_tree_node));

	// BUG_ON(radix_tree_is_indirect_ptr(ret));
	return ret;
}

static int radix_tree_extend(struct radix_tree_root *root, unsigned long index)
{
	struct radix_tree_node *node;
	unsigned int height;
	int tag;

	printf("this is %s(): %d\r\n", __func__, __LINE__);

	/* Figure out what the height should be.  */
	height = root->height + 1;
	printf("this is %s(): %d >>> height = %d, root->height = %d\r\n", __func__, __LINE__, height, root->height);
	while (index > radix_tree_maxindex(height))
		height++;

	printf("this is %s(): %d >>> index = %d, radix_tree_maxindex(height) = %d\r\n", __func__, __LINE__, index, radix_tree_maxindex(height));
	printf("this is %s(): %d >>> height = %d, root->height = %d\r\n", __func__, __LINE__, height, root->height);

	if (root->rnode == NULL) {
		printf("this is %s(): %d\r\n", __func__, __LINE__);
		root->height = height;
		goto out;
	}

	printf("this is %s(): %d\r\n", __func__, __LINE__);
	do {
		printf("this is %s(): %d >>> height = %d, root->height = %d\r\n", __func__, __LINE__, height, root->height);
		unsigned int newheight;
		if (!(node = radix_tree_node_alloc(root)))
			return -ENOMEM;

		/* Increase the height.  */
		node->slots[0] = radix_tree_indirect_to_ptr(root->rnode);

		/* Propagate the aggregated tag info into the new root */
		for (tag = 0; tag < RADIX_TREE_MAX_TAGS; tag++) {
			printf("this is %s(): %d\r\n", __func__, __LINE__);
			if (root_tag_get(root, tag))
				tag_set(node, tag, 0);
		}

		newheight = root->height+1;
		node->height = newheight;
		node->count = 1;
		node = radix_tree_ptr_to_indirect(node);
		rcu_assign_pointer(root->rnode, node);
		root->height = newheight;
		printf("this is %s(): %d\r\n", __func__, __LINE__);
	} while (height > root->height);
out:
	printf("this is %s(): %d\r\n", __func__, __LINE__);
	return 0;
}

int radix_tree_insert(struct radix_tree_root *root,
			unsigned long index, void *item)
{
	struct radix_tree_node *node = NULL, *slot;
	unsigned int height, shift;
	int offset;
	int error;

	printf("this is %s(): %d >>> index = %d\r\n", __func__, __LINE__, index);
	printf("this is %s(): %d >>> root->height = %d\r\n", __func__, __LINE__, root->height);
	printf("this is %s(): %d >>> radix_tree_maxindex(root->height) = %d\r\n", __func__, __LINE__, radix_tree_maxindex(root->height));

	// BUG_ON(radix_tree_is_indirect_ptr(item));

	/* Make sure the tree is high enough.  */
	if (index > radix_tree_maxindex(root->height)) {
		printf("this is %s(): %d\r\n", __func__, __LINE__);
		error = radix_tree_extend(root, index);
		if (error)
			return error;
	}

	printf("this is %s(): %d\r\n", __func__, __LINE__);
	slot = radix_tree_indirect_to_ptr(root->rnode);
	printf("this is %s(): %d\r\n", __func__, __LINE__);

	height = root->height;
	shift = (height-1) * RADIX_TREE_MAP_SHIFT;

	offset = 0;			/* uninitialised var warning */
	while (height > 0) {
		printf("this is %s(): %d >>> height = %d\r\n", __func__, __LINE__, height);
		if (slot == NULL) {
			printf("this is %s(): %d\r\n", __func__, __LINE__);
			/* Have to add a child node.  */
			if (!(slot = radix_tree_node_alloc(root)))
				return -ENOMEM;
			slot->height = height;
			if (node) {
				rcu_assign_pointer(node->slots[offset], slot);
				node->count++;
			} else
				rcu_assign_pointer(root->rnode,
					radix_tree_ptr_to_indirect(slot));
		}

		/* Go a level down */
		offset = (index >> shift) & RADIX_TREE_MAP_MASK;
		node = slot;
		slot = node->slots[offset];
		shift -= RADIX_TREE_MAP_SHIFT;
		height--;
		printf("this is %s(): %d\r\n", __func__, __LINE__);
	}
	printf("this is %s(): %d\r\n", __func__, __LINE__);

	if (slot != NULL)
		return -EEXIST;

	printf("this is %s(): %d\r\n", __func__, __LINE__);
	if (node) {
		printf("this is %s(): %d\r\n", __func__, __LINE__);
		node->count++;
		rcu_assign_pointer(node->slots[offset], item);
		// BUG_ON(tag_get(node, 0, offset));
		// BUG_ON(tag_get(node, 1, offset));
	} else {
		printf("this is %s(): %d\r\n", __func__, __LINE__);
		rcu_assign_pointer(root->rnode, item);
		// BUG_ON(root_tag_get(root, 0));
		// BUG_ON(root_tag_get(root, 1));
	}

	printf("this is %s(): %d\r\n", __func__, __LINE__);
	return 0;
}

static void *radix_tree_lookup_element(struct radix_tree_root *root,
				unsigned long index, int is_slot)
{
	unsigned int height, shift;
	struct radix_tree_node *node, **slot;

	node = rcu_dereference(root->rnode);
	if (node == NULL)
		return NULL;

	if (!radix_tree_is_indirect_ptr(node)) {
		if (index > 0)
			return NULL;
		return is_slot ? (void *)&root->rnode : node;
	}
	node = radix_tree_indirect_to_ptr(node);

	height = node->height;
	if (index > radix_tree_maxindex(height))
		return NULL;

	shift = (height-1) * RADIX_TREE_MAP_SHIFT;

	do {
		slot = (struct radix_tree_node **)
			(node->slots + ((index>>shift) & RADIX_TREE_MAP_MASK));
		node = rcu_dereference(*slot);
		if (node == NULL)
			return NULL;

		shift -= RADIX_TREE_MAP_SHIFT;
		height--;
	} while (height > 0);

	return is_slot ? (void *)slot:node;
}

void *radix_tree_lookup(struct radix_tree_root *root, unsigned long index)
{
	return radix_tree_lookup_element(root, index, 0);
}

void **radix_tree_lookup_slot(struct radix_tree_root *root, unsigned long index)
{
	return (void **)radix_tree_lookup_element(root, index, 1);
}

static unsigned long __maxindex(unsigned int height)
{
	unsigned int width = height * RADIX_TREE_MAP_SHIFT;
	int shift = RADIX_TREE_INDEX_BITS - width;

	if (shift < 0)
		return ~0UL;
	if (shift >= BITS_PER_LONG)
		return 0UL;
	return ~0UL >> shift;
}

static void radix_tree_init_maxindex(void)
{
	unsigned int i;

	for (i = 0; i < ARRAY_SIZE(height_to_maxindex); i++) {
		height_to_maxindex[i] = __maxindex(i);
		printf("this is %s(): %d >>> height_to_maxindex[%d] = %d\r\n", __func__, __LINE__, i, height_to_maxindex[i]);
	}
}

void radix_tree_init(void)
{
	// radix_tree_node_cachep = kmem_cache_create("radix_tree_node",
	// 		sizeof(struct radix_tree_node), 0,
	// 		SLAB_PANIC | SLAB_RECLAIM_ACCOUNT,
	// 		radix_tree_node_ctor);
	radix_tree_init_maxindex();
	// hotcpu_notifier(radix_tree_callback, 0);
}
